package jsr292.weaver.opt;

import java.io.PrintWriter;
import java.lang.instrument.ClassDefinition;
import java.lang.instrument.UnmodifiableClassException;
import java.lang.reflect.Field;
import java.lang.reflect.Modifier;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

import jsr292.java.lang.invoke.BootstrapCallSite;
import jsr292.java.lang.invoke.CallSite;
import jsr292.java.lang.invoke.MethodHandle;
import jsr292.weaver.Agent;
import jsr292.weaver.Capabilities;
import jsr292.weaver.Debug;
import jsr292.weaver.opt.OptimizerMethodInfo.Kind;

import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.util.CheckClassAdapter;


public class Optimizer {
    // OPTIMIZER_*_THRESHOLD must be greater than OPTIMIZER_*_OPTIMIZE_THRESHOLD
    // all values must be greater than zero
    public static final int OPTIMIZER_MH_THRESHOLD = 1000;
    public static final int OPTIMIZER_MH_OPTIMIZE_THRESHOLD = 500;
    public static final int OPTIMIZER_INDY_THRESHOLD = 100;
    public static final int OPTIMIZER_INDY_OPTIMIZE_THRESHOLD = 50;
    public static final int OPTIMIZER_INDY_DEOPT_THRESHOLD = OPTIMIZER_INDY_OPTIMIZE_THRESHOLD / 2;
    
    private static AtomicInteger boundMapCounter = new AtomicInteger();
    private static final ConcurrentHashMap<String, byte[]> classDefinitionMap=
	new ConcurrentHashMap<String, byte[]>();
	
    /** Called by the agent
     */	
    public static void registerClassDefinition(String internalClassName, byte[] bytecode) {
        if (!Capabilities.canRetransform()) {
	  classDefinitionMap.put(internalClassName.replace('/', '.'), bytecode);
        }
    }
    
    private static void registerClassDefinition(Class<?> clazz, byte[] bytecode) {
        if (!Capabilities.canRetransform()) {
	  classDefinitionMap.put(clazz.getName(), bytecode);
        }
    }
    
    private static ClassDefinition getClassDefinition(Class<?> dynClass) {
	String name = dynClass.getName();
        byte[] bytecode = classDefinitionMap.get(name);
        if (bytecode == null)
            return null;
	return new ClassDefinition(dynClass, bytecode);
    }
    
    public static int nextBoundMapId() {
	return boundMapCounter.getAndIncrement();
    }
    
    
    /** Entry point called by generated code.
     */
    public static void retransform(Class<?> clazz) {
        synchronized(BootstrapCallSite.getCallSiteGlobalLock()) {
            retransform(Collections.singleton(clazz));
        }
    }
    
    /** Entry point called by generated code.
     */
    public static void deoptimize(Class<?> clazz) {
        StackTraceElement[] stackTrace = new Throwable().getStackTrace();
        if (stackTrace.length < 2) {
            // here, we are in trouble because we have no idea where we come from
            // so we do nothing, XXX may revisit this idea
            return;
        }
        
        StackTraceElement callerElement = stackTrace[1];
        String callerMethodName = callerElement.getMethodName();
        //System.out.println("deoptimize from caller stub "+clazz.getName()+"."+callerMethodName);
        
        // set invokedynamic/method handle counter to ask re-profiling
        try {
            if (callerMethodName.startsWith("$method_handle_stub")) {
                // it's a method handle stub
                Field field = clazz.getField(callerMethodName+"counter");
                field.setInt(null, 0);
            } else
                if (callerMethodName.startsWith("$indy_stub")) {
                    // it's an invokedynamic stub
                    Field field = clazz.getField(callerMethodName+"cs");
                    CallSite callSite = (CallSite)field.get(null);
                    callSite.counter = 0;
                } else {
                    throw new AssertionError("deopt must be called from an indy/mh stub");
                }
        } catch (NoSuchFieldException e) {
            throw (AssertionError)new AssertionError().initCause(e);
        } catch (IllegalAccessException e) {
            throw (AssertionError)new AssertionError().initCause(e);
        }
        
        synchronized(BootstrapCallSite.getCallSiteGlobalLock()) {
            retransform(Collections.singleton(clazz));
        }
    }
    
    // this method must only be called if the current thread held
    // the CallSite global lock to create a global de-reoptimization order
    public static void retransform(Set<? extends Class<?>> classes) {
        //System.err.println("ask instrumentation retransform "+classes);
        try {
            if (!Capabilities.canRetransform()) {   // path compatible with JDK5
                ArrayList<ClassDefinition> classDefinitions = new ArrayList<ClassDefinition>(classes.size());
                for(Class<?> clazz: classes) {
                    ClassDefinition classDefinition = getClassDefinition(clazz);
                    if (classDefinition != null) {
                        classDefinitions.add(classDefinition);
                    }
                }
                if (classDefinitions.isEmpty())
                    return;
                try {
                    Agent.instrumentation.redefineClasses(
                            classDefinitions.toArray(new ClassDefinition[classDefinitions.size()]));
                } catch (ClassNotFoundException e) {
                    // should not be called
                    throw new AssertionError(e);
                }
            } else {                                // path compatible with JDK6
                Capabilities.retransform(classes);
            }
        } catch (UnmodifiableClassException e) {
            throw new AssertionError(e);
        }
    }

    public static void debugCounter(Class<?> dynClass) {
	System.err.println("-- debug counter --");
	
	Field[] fields = dynClass.getFields();
	for(Field field:fields) {
	    if (!Modifier.isStatic(field.getModifiers()))
		continue;
	    
	    String fieldName = field.getName();
	    if (fieldName.startsWith("$method_handle_stub") && fieldName.endsWith("$counter")) {
		Object counter;
		try {
		    counter = field.get(null);
		} catch (IllegalAccessException e) {
		    continue;
		}
		System.err.println("method handle counter "+fieldName+" "+counter);
		continue;
	    }
	    if (fieldName.startsWith("$indy_stub") && fieldName.endsWith("$cs")) {
		CallSite callsite;
		try {
		    callsite = (CallSite)field.get(null);
		} catch (IllegalAccessException e) {
		    continue;
		}
		System.err.println("call site counter "+fieldName+" "+callsite.counter);
		continue;
	    }   
	}
    }
    
    // Called by the agent
    public static byte[] reweave(Class<?> classBeingRedefined, byte[] classfileBuffer) {
        //System.err.println("reweave "+classBeingRedefined);
        
        Field[] fields;
	try {
	    fields = classBeingRedefined.getFields();
	} catch (SecurityException e) {
	    throw new AssertionError(e);
	}
	
	HashMap<String, OptimizerMethodInfo> methodInfoMap = new HashMap<String, OptimizerMethodInfo>();
	for(Field field:fields) {
	    int modifiers = field.getModifiers();
	    if (!Modifier.isStatic(modifiers) || !Modifier.isPublic(modifiers)) {
		continue;
	    }
	    
	    String fieldName = field.getName();
	    if (field.getType()==int.class &&
		fieldName.startsWith("$method_handle_stub") &&
		fieldName.endsWith("$counter")) {
		handleMethodHandleStub(classBeingRedefined, field, methodInfoMap);
	    } else {
		if (field.getType()==CallSite.class &&
		    fieldName.startsWith("$indy_stub") &&
		    fieldName.endsWith("$cs")) {
		    handleIndyStub(classBeingRedefined, field, methodInfoMap);
		}
	    }
	}
	
	if (Debug.DEBUG) {
	  System.out.println("methodInfo map "+methodInfoMap);
	}
	
	if (methodInfoMap.isEmpty()) {
	    return null;
	}
	
	ClassReader reader = new ClassReader(classfileBuffer);
        ClassWriter writer = new ClassWriter(reader, ClassWriter.COMPUTE_MAXS);
        
        reader.accept(new OptimizerWeaver(writer, classBeingRedefined, methodInfoMap), ClassReader.SKIP_FRAMES);
        byte[] byteArray = writer.toByteArray();
        
        registerClassDefinition(classBeingRedefined, byteArray);
        return byteArray;
    }
    
    private static void ensureVisibility(Field field) {
        if (Modifier.isPublic(field.getDeclaringClass().getModifiers()))
            return;
        try {
            field.setAccessible(true);
        } catch (SecurityException e) {
            throw new AssertionError(e);
        }
    }
    
    private static void handleMethodHandleStub(Class<?> classBeingRedefined, Field field, HashMap<String, OptimizerMethodInfo> methodInfoMap) {
        ensureVisibility(field); // field can be own by a non public class
            
        int counter;
        try {
	    counter = field.getInt(null);

	    // if method handle not called enough, no need to optimize
	    if (counter >= 0 && counter < OPTIMIZER_MH_OPTIMIZE_THRESHOLD) {
		return;
	    }	
	} catch(SecurityException e) {
	    throw new AssertionError(e);
	} catch (IllegalAccessException e) {
	    throw new AssertionError(e);
	}

	String fieldName = field.getName();
	int index = fieldName.lastIndexOf('$');
	String stubName = fieldName.substring(0, index+1);
	Field handleField;
	try {
	    handleField = classBeingRedefined.getField(stubName+ "mh");
	} catch (SecurityException e) {
	    throw new AssertionError(e);
	} catch (NoSuchFieldException e) {
	    throw new AssertionError(e);
	}

	if (handleField.getType() != MethodHandle.class) {
	    throw new AssertionError("method handle slot field doesn't store a method handle");
	}
	
	if (Debug.DEBUG) {
            System.out.println("Stub site "+field.getName()+" name "+ stubName+" "+" counter "+counter);
        }
	
	ensureVisibility(handleField);
	
	MethodHandle mh;
	try {
	    mh = (MethodHandle)handleField.get(null);
	} catch (IllegalAccessException e) {
	    throw new AssertionError(e);
	} catch(SecurityException e) {
	    return;
	}
	
	try {
	    // avoid cyclic optimizations
	    field.set(null, Integer.MIN_VALUE);
	} catch (IllegalAccessException e) {
	    throw new AssertionError(e);
	} catch(SecurityException e) {
	    throw new AssertionError(e);
	}
	
	// stupid case, but if a user uses a null method handle,
        // NPE should not be thrown here but will be throw later when doing the call 
        if (mh == null) {
            return;
        }
	
        if (!mh.isWeavable(classBeingRedefined)) {
	    if (Debug.DEBUG) {
                System.err.println("stub site "+field.getName()+" name "+stubName+" is not weavable !");
                System.err.println("  mh: "+mh);
            }
	    return;
	}

	methodInfoMap.put(stubName, new OptimizerMethodInfo(Kind.METHOD_HANDLE_INVOKE, mh, null));
    }
    
    private static void handleIndyStub(Class<?> classBeingRedefined, Field field, HashMap<String, OptimizerMethodInfo> methodInfoMap) {
        ensureVisibility(field);  // field can be own by a non public class
        
        CallSite callSite;
	try {
	    callSite = (CallSite)field.get(null);
	} catch (IllegalAccessException e) {
	    throw new AssertionError(e);
	} catch(SecurityException e) {
	    return;
	}
	
	int counter = callSite.counter;
        if (Debug.DEBUG) {
          System.out.println("Call site "+field.getName()+" callSite "+ callSite+" "+" counter "+counter);
	}
	
        // if callsite haven't been called enough, don't optimize
	if (counter >= 0 && counter < OPTIMIZER_INDY_OPTIMIZE_THRESHOLD) {
	    return;
	}
	
	// avoid cyclic optimizations
	callSite.counter = Integer.MIN_VALUE;
	
	MethodHandle mh = callSite.target;
	if (!mh.isWeavable(classBeingRedefined)) {
	    if (Debug.DEBUG) {
	        System.err.println("Call site "+field.getName()+" callSite "+callSite+" is not weavable !");
	        System.err.println("  mh: "+mh);
	    }
	    return;
	}
	
	String fieldName = field.getName();
	int index = fieldName.lastIndexOf('$');
	String stubName = fieldName.substring(0, index+1);
	methodInfoMap.put(stubName, new OptimizerMethodInfo(Kind.INDY_INVOKE, mh, callSite));
    }
}
